import requests
import json
import re
import base64
from datetime import datetime, timedelta

# ========== DeepSeek 配置 ==========
API_KEY = "sk-d2513b4c4626409599a73ba64b2e9033" 
API_URL = "https://api.deepseek.com/v1/chat/completions"
HEADERS = {
    "Content-Type": "application/json",
    "Authorization": f"Bearer {API_KEY}"
}


def encode_image_to_base64(filepath):
    try:
        with open(filepath, "rb") as f:
            return base64.b64encode(f.read()).decode("utf-8")
    except FileNotFoundError:
        print(f"[错误] 图片路径不存在: {filepath}")
        return ""

def millis_to_iso(ts):
    return datetime.utcfromtimestamp(int(ts) / 1000).isoformat() + "Z"

def extract_json_from_response(response_text):
    match = re.search(r"```json\s*(.*?)\s*```", response_text, re.DOTALL)
    json_str = match.group(1) if match else response_text
    try:
        return json.loads(json_str)
    except json.JSONDecodeError as e:
        print("无法解析 JSON：", e)
        return None

def filter_and_format_messages(raw_messages):
    filtered = []
    for msg in raw_messages:
        if msg["msg_type"] == "text":
            text = msg["message_content"]["text"]
            if not any(x in text for x in ["收到", "谢谢", "了解了"]):
                filtered.append({
                    "msg_type": "text",
                    "message_content": {
                        "text": text
                    },
                    "create_time": millis_to_iso(msg["create_time"])
                })
        elif msg["msg_type"] == "image":
            filtered.append({
                "msg_type": "image",
                "message_content": {
                    "image_key": msg["message_content"]["image_key"],
                    "base64": msg.get("base64", "")
                },
                "create_time": millis_to_iso(msg["create_time"])
            })
    return filtered

def cluster_and_classify(messages):
    prompt_lines = [
        f"- [{msg['create_time']}] [TEXT] {msg['message_content']['text']}"
        for msg in messages if msg["msg_type"] == "text"
    ]

    prompt = f"""
<context>
你是一个建筑工地巡检助手。在巡检过程中，你会将工地现场发送信息留言中的问题进行聚类与合理按标签分类。
同时你会考虑聚类准确性，以及合理分类。你会预先规划、使用最少步骤完成任务，并在遇到不确定性时请求澄清。
</context>

<planning_rules>

你需要从以下巡检留言中：

1. 聚类出多个事件，每个事件卡对应一个或多个留言，其事件卡内容包括：
    - 一个简洁的 summary；
    - 所属的原始留言 messages（保留原始字段，如 msg_type, message_content 和 create_time）；
2. 对每个事件卡进行分类，类别为以下五类之一：
    - 巡检问题：发现现场隐患或风险点
    - 验收：施工完成后的检查
    - 旁站：施工过程中的现场陪同监督
    - 闭环：对之前问题的处理、确认
    - 其他：不属于以上类型的其他事项

</planning_rules>

<format_rules>
请以以下 JSON 格式返回所有事件卡：

[
  {{
    "event_id": 1,
    "summary": "...",
    "messages": [
      {{
        "msg_type": "text",
        "message_content": {{ "text": "..." }},
        "create_time": "..."
      }},
      ...
    ],
    "class": "巡检问题"
  }},
  ...
]
</format_rules>

以下是巡检留言信息：
{chr(10).join(prompt_lines)}
"""

    payload = {
        "model": "deepseek-chat",
        "messages": [
            {"role": "system", "content": "你是一个结构化事件聚类与分类助手"},
            {"role": "user", "content": prompt}
        ],
        "temperature": 0.3
    }

    try:
        response = requests.post(API_URL, headers=HEADERS, json=payload)
        response.raise_for_status()
        return extract_json_from_response(response.json()["choices"][0]["message"]["content"])
    except requests.exceptions.RequestException as e:
        print("请求失败:", e)
        return None

def assign_images_to_events(events, images, time_window_minutes=10):

    for event in events:

        times = [
            datetime.fromisoformat(msg["create_time"].replace("Z", "+00:00"))
            for msg in event["messages"] if msg["msg_type"] == "text"
        ]
        window_start = min(times) - timedelta(minutes=time_window_minutes)
        window_end = max(times) + timedelta(minutes=time_window_minutes)
        event["candidate_images"] = [
            {
                "message_content": {
                    "image_key": img["message_content"]["image_key"],
                    "base64": img["message_content"]["base64"]
                },
                "create_time": img["create_time"]
            }
            for img in images if window_start <= datetime.fromisoformat(img["create_time"].replace("Z", "+00:00")) <= window_end
        ]
    return events


tools = [
    {
        "type": "function",
        "function": {
            "name": "cluster_and_classify_events",
            "description": "将工地巡检的文本消息聚类成事件，并为每个事件进行分类",
            "parameters": {
                "type": "object",
                "properties": {
                    "messages": {
                        "type": "array",
                        "description": "输入的文本消息数组",
                        "items": {
                            "type": "object",
                            "properties": {
                                "msg_type": {"type": "string"},
                                "message_content": {"type": "object"},
                                "create_time": {"type": "string"}
                            },
                            "required": ["msg_type", "message_content", "create_time"]
                        }
                    }
                },
                "required": ["messages"]
            }
        }
    }
]

tool_name = [tool["function"]["name"] for tool in tools]
print(f"创建了{len(tools)}个工具，为：{tool_name}\n")

messages = [
    {
        "role": "system",
        "content": "你是一个建筑现场信息分析专家，可以聚类和分类文本信息。"
    },
    {
        "role": "user",
        "content": "请帮我对以下巡检信息进行处理分析。如果需要聚类和分类文本信息，请调用 `cluster_and_classify_events` 函数。"
    }
]
print("messages 数组创建完成\n")

def function_call(function):
    payload = {
        "model": "deepseek-chat",
        "messages": messages,
        "tools": tools,
        "tool_choice": {
            "type": "function",
            "function": {
                "name": "cluster_and_classify_events"
            }
        },
        "temperature": 0.3
    }

    try:
        response = requests.post(API_URL, headers=HEADERS, json=payload)
        response.raise_for_status()
        call_result = response.json()
        tool_response_text = call_result["choices"][0]["message"]["tool_calls"][0]["function"]["arguments"]
        return json.loads(tool_response_text)
    except Exception as e:
        print("Function Call 请求失败:", e)
        return None

def main():
    img2 = encode_image_to_base64("IMG_2.jpg")
    img3 = encode_image_to_base64("image03.jpg")

    raw_messages = [
        {
            "message_id": "om_3",
            "msg_type": "text",
            "create_time": "1711933800000",  # 2024-03-31 22:30:00 UTC
            "sender_id": "user1",
            "message_content": {"text": "在3号楼东南角发现脚手架松动，有坠落风险"}
        },
        {
            "message_id": "om_5",
            "msg_type": "text",
            "create_time": "1711933870000",  # 2024-03-31 22:31:10 UTC
            "sender_id": "user1",
            "message_content": {"text": "责任人张三已通知整改"}
        },
        {
            "message_id": "om_6",
            "msg_type": "text",
            "create_time": "1711934100000",  # 2024-03-31 22:35:00 UTC
            "sender_id": "user1",
            "message_content": {"text": "收到，谢谢"}
        },
        {
            "message_id": "om_7",
            "msg_type": "text",
            "create_time": "1711940400000",  # 2024-03-31 23:00:00 UTC
            "sender_id": "user1",
            "message_content": {"text": "2号楼发现电缆裸露，存在触电危险"}
        },
        {
            "message_id": "om_8",
            "msg_type": "image",
            "create_time": "1711940460000",  # 2024-03-31 23:01:00 UTC
            "sender_id": "user2",
            "message_content": {"image_key": "IMG_2.jpg"},
            "base64": img2
        },
        {
            "message_id": "om_9",
            "msg_type": "text",
            "create_time": "1711940700000",  # 2024-03-31 23:05:00 UTC
            "sender_id": "user2",
            "message_content": {"text": "监理回复：了解了"}
        },
        {
            "message_id": "om_extra_07",
            "msg_type": "text",
            "create_time": "1711952400000",  # 2024-04-01 02:00:00 UTC
            "sender_id": "user5",
            "message_content": {"text": "今日上午9点，3号楼主体钢筋绑扎作业开始"}
        },
        {
            "message_id": "om_extra_08",
            "msg_type": "text",
            "create_time": "1711953300000",  # 2024-04-01 02:15:00 UTC
            "sender_id": "user5",
            "message_content": {"text": "监理张工已到场确认施工规范"}
        },
        {
            "message_id": "om_extra_05",
            "msg_type": "text",
            "create_time": "1712028300000",  # 2024-04-02 08:45:00 UTC
            "sender_id": "user4",
            "message_content": {"text": "5号楼地下室防水层施工于昨日完成，今日进行初步验收"}
        },
        {
            "message_id": "om_extra_06",
            "msg_type": "text",
            "create_time": "1712028600000",  # 2024-04-02 08:50:00 UTC
            "sender_id": "user4",
            "message_content": {"text": "验收人员已到场，准备进行验收"}
        },
        {
            "message_id": "om_a",
            "msg_type": "image",
            "create_time": "1712029000000",  
            "sender_id": "user4",
            "message_content": {"image_key": "image03.jpg"},
            "base64": img3
        }
    ]

    filtered = filter_and_format_messages(raw_messages)
    text_msgs = [m for m in filtered if m["msg_type"] == "text"]
    image_msgs = [m for m in filtered if m["msg_type"] == "image"]

    print("\n所有图像消息:")
    for im in image_msgs:
        print(f"- {im['message_content']['image_key']} @ {im['create_time']}")


    events = cluster_and_classify(text_msgs)

    events = assign_images_to_events(events, image_msgs)

    print("\n最终事件卡：")
    for e in events:
        print(f"事件 {e['event_id']}: [{e['class']}] - {e['summary']}")
        print("  Messages:")
        for msg in e['messages']:
            msg_type = msg['msg_type']
            content = msg['message_content'].get("text", msg['message_content'].get("image_key", ""))
            timestamp = msg['create_time']
            print(f"    - [{msg_type}] {timestamp} | {content}")
        print(f"  Candidate Images: {len(e.get('candidate_images', []))}")


if __name__ == "__main__":
    main()
